#!/usr/bin/env python3
# -*- coding: utf-8 -*-
'''
计算探头夹角程序
'''
import numpy as np
import matplotlib.pyplot as plt
import sys
import os
import math
import astropy.coordinates as coord
from astropy.coordinates import SkyCoord
from astropy import units as u
import Quaternion as Qt
from astropy.io import fits
from astropy.time import Time
from astropy.coordinates.representation import (CartesianRepresentation, UnitSphericalRepresentation)
from astropy import coordinates


def Sat_J2000(theta=[0.0], phi=[0.0], q1s=[-0.309514, -0.330917], q2s=[-0.656108, -0.624895],
              q3s=[0.351349, 0.330917], q0s=[0.591842, 0.624895]):
    # theta,phi in degree. list or array
    # time format: yyyy-mm-ddTHH:MM:SS
    RADECS = []
    q1s = np.array(q1s)
    q2s = np.array(q2s)
    q3s = np.array(q3s)
    q0s = np.array(q0s)
    theta = np.array(theta)
    phi = np.array(phi)
    ra = np.empty(q1s.shape, dtype=q1s.dtype)
    dec = np.empty(q1s.shape, dtype=q1s.dtype)
    q_xyz2zxy = Qt.Quat(([[1, 0, 0], [0, -1, 0], [0, 0, -1]]))
    #vstack按行堆叠成一个数组
    quat = np.vstack((q1s, q2s, q3s, q0s))
    #enumerate同时获取索引和值
    for index, quat1 in enumerate(quat.T):
        quat1 = Qt.Quat(Qt.normalize(quat1))
        #        quat1 = q_xyz2zxy * quat1
        quat1 = quat1 * q_xyz2zxy
        sph1 = UnitSphericalRepresentation(phi[index] * u.deg, (90.0 - theta[index]) * u.deg)
        cart1 = sph1.represent_as(CartesianRepresentation)
        cart1_xyz = cart1.xyz
        cart2_xyz = np.matmul(quat1.transform, cart1_xyz)  #np.matmu 两个矩阵的乘积
        cart2 = CartesianRepresentation(cart2_xyz[0], cart2_xyz[1], cart2_xyz[2])
        sph2 = cart2.represent_as(UnitSphericalRepresentation)

        ra[index] = sph2.lon.deg
        dec[index] = sph2.lat.deg
    return [ra, dec]


#         print('RA Dec: ', ra[index], dec[index])
#         print('theta phi: ',theta[index],phi[index])
def J2000_Sat(ra=[180 + 5.4403931, 0], dec=[-3.01239707, 0], q1s=[-0.309514, -0.330917], q2s=[-0.656108, -0.624895],
              q3s=[0.351349, 0.330917], q0s=[0.591842, 0.624895]):
    # ra,dec in degree. list or array
    # time format: yyyy-mm-ddTHH:MM:SS
    ra = np.array(ra)
    dec = np.array(dec)
    q1s = np.array(q1s)
    q2s = np.array(q2s)
    q3s = np.array(q3s)
    q0s = np.array(q0s)
    quat = np.vstack((q1s, q2s, q3s, q0s))
    theta = np.empty(q1s.shape, dtype=q1s.dtype)
    phi = np.empty(q1s.shape, dtype=q1s.dtype)
    q_xyz2zxy = Qt.Quat(([[1, 0, 0], [0, -1, 0], [0, 0, -1]]))
    for index, quat1 in enumerate(quat.T):
        quat1 = Qt.Quat(Qt.normalize(quat1))
        quat1 = quat1 * q_xyz2zxy
        sph1 = UnitSphericalRepresentation(ra[index] * u.deg, dec[index] * u.deg)
        cart1 = sph1.represent_as(CartesianRepresentation)  # 将RADEC转换为xyz
        cart1_xyz = cart1.xyz
        cart2_xyz = np.matmul(quat1.transform.T, cart1_xyz)  # matmul矩阵相乘
        cart2 = CartesianRepresentation(cart2_xyz[0], cart2_xyz[1], cart2_xyz[2])
        sph2 = cart2.represent_as(UnitSphericalRepresentation)

        phi[index] = sph2.lon.deg
        theta[index] = 90 - sph2.lat.deg
        phitheta = [[theta[i], phi[i]] for i in range(len(phi))]
    return phitheta

def utc2met(T):
    return (Time(T, format='isot', scale='utc') - Time('2019-01-01T00:00:00.000', format='isot', scale='utc')).to(
        u.s).value

def Quaternion_get(location_file_posatt, TriggerTime):
    dulist_posatt = fits.open(location_file_posatt)
    Flag = np.array(dulist_posatt[1].data, dtype=object)
    if (Flag.size == 0):
        print('posatt file is exit, but is empty!')
    else:
        Time_GECAM = np.round(dulist_posatt[1].data.field('TIME'))
        # print(min(Time_GECAM),max(Time_GECAM))
        Q1 = dulist_posatt[1].data.field('Q1')
        Q2 = dulist_posatt[1].data.field('Q2')
        Q3 = dulist_posatt[1].data.field('Q3')
        Num_GECAM = np.where(Time_GECAM == int(TriggerTime))
        if (len(Num_GECAM[0]) > 0):
            q3s = [Q1[Num_GECAM[0]][0], Q2[Num_GECAM[0]][0], Q3[Num_GECAM[0]][0]]
        else:
            q3s = [0,0,0]
            print('No this time Quaternion!')
    return q3s

def angle(location_file_posatt, utc_input,ra,dec):
    ra = [float(ra)]
    dec = [float(dec)]
    TriggerTime = utc2met(utc_input)
    q3s = Quaternion_get(location_file_posatt, TriggerTime)
    phitheta=J2000_Sat(ra=ra, dec=dec, q1s=[q3s[0] for i in range(len(ra))], q2s=[q3s[1] for i in range(len(ra))],
              q3s=[q3s[2] for i in range(len(ra))], q0s=[np.sqrt(1 - sum(np.array(q3s) ** 2)) for i in range(len(ra))])
    a = [0, 40, 40, 40, 40, 40, 40, 73.5, 73.5, 73.5, 73.5, 73.5, 73.5, 73.5, 73.5, 73.5, 73.5, 90, 90, 90, 90, 90, 90,
         90, 90]
    b = np.array(
        [0.00, 242.17, 188.67, 135.17, 62.17, 8.67, 315.17, 260.50, 224.50, 188.50, 152.50, 116.50, 80.50, 44.50, 8.50,
         332.50, 296.50, 270.00, 215.00,
         180.00, 125.00, 90.00, 35.00, 0.00, 305.00]) + 90
    for i in range(len(b)):
        if b[i] >= 360:
            b[i] = b[i] - 360
        b[0] = 0

    A = a
    B = list(b)
    XL = []
    #np.radians 将角度转为弧度
    for k in range(len(A)):
        XL.append(np.array(
            [np.sin(np.radians(A[k])) * np.cos(np.radians(B[k])), np.sin(np.radians(A[k])) * np.sin(np.radians(B[k])),
             np.cos(np.radians(A[k]))]))

    ###计算一个入射角度和各个GRD的夹角,输入thetaphi
    GRB = phitheta[0]
    GRB = np.array([np.sin(np.radians(GRB[0])) * np.cos(np.radians(GRB[1])),
                    np.sin(np.radians(GRB[0])) * np.sin(np.radians(GRB[1])),
                    np.cos(np.radians(GRB[0]))])
    JIAJIAO = {}
    for i in range(len(A)):
        JIAJIAO['GRD%d'%(i+1)] = np.arccos(round(GRB.dot(XL[i]), 5)) * 180 / np.pi
    print('源与各GRD探头的夹角为:')
    print(JIAJIAO)

if __name__ == '__main__':
    location_file_posatt = 'C:/Users/Dell/Desktop/GECAM/posatt/gb_posatt_tn221014_192139_v00.fits'
    # utc_input = input()
    utc_input ='2022-10-14T19:21:39.050'
    # radec = input()
    # ra = [float(radec.split(',')[0])]
    # dec = [float(radec.split(',')[1])]
    ra = 293.72
    dec = 21.88
    #q3s = [-0.4305435, -0.8457335, 0.3111986]
    angle(location_file_posatt, utc_input,ra,dec)


